#!/home/zhonc0a/Enthought/Canopy_64bit/User/bin/python
import math
import re
import argparse
import os
import sys
import pylab
class dihedralScan():
    def __init__(self,filename):
        self.filename=filename
        self.coords=[]   
        self.relativeE=[]  
        self.energy=[]
        self.energy=[]
        self.unit='kcal'
    def readLog(self):
        logFile=open(self.filename)
        step=None
        energy=None
        se=[]
        scanCoord=None
        coords=[]
        
        for line in logFile:
            if 'SCF Done' in line:
                energy=line.split()[4]
            if 'Step number' in line:
                step=line.split()[12]
                se.append((step,energy))
            if scanCoord:
                if scanCoord in line:
                    coords.append(abs(float(line.split()[3])))
            if re.search('!.*Scan',line):
                scanCoord=line.split()[2]
        self.coords=coords
        oldStep=None
        newStep=None
        finalSE=[]            
        for i in se:
            newStep=i
            if oldStep:
                if oldStep[0] != newStep[0]:
                    finalSE.append(oldStep)
            oldStep=newStep
        finalSE.append(se[-1])
        minEnergy=min([float(i[1]) for i in finalSE])
        self.relativeE=[(float(i[1])-minEnergy) for i in finalSE]
        if self.unit=='au':
            self.energy=self.relativeE
        if self.unit=='KCal':  
            self.energy=[i*627.509469 for i in self.relativeE]
        if self.unit=='eV':
            self.energy=[i*27.211385 for i in self.relativeE]


parser=argparse.ArgumentParser(description='extract energy and angle from Gaussian09 output file (relax scan) and draw PES')
parser.add_argument("-u","--unit",help='unit of energy(KCal,eV, or au). Default: kcal',default='KCal')
parser.add_argument("-y","--ymax",help='maximum value of y axis. Default is the maximum value * 1.1')
parser.add_argument("-d","--dpi",help='set the dpi of the figure. Default is 150',type=int,default=150)
parser.add_argument("-s","--size",help='set the size of the figure(two number seperate by comma, unit is inch), e.g. -s (8,6).')
parser.add_argument("-p","--pic",help='save the pic with format (svg,png,ps,eps,pdf), default do not save pic',default=None)
parser.add_argument("-t","--tick",help='set the increment of ticks on specific axis (the formats are [hvo][xy][num],[hvo][xy][num], e.g. hy1,vx30 means ticks on y axis increase by 1 in horizontal graph and ticks on x axis increase by 30 in vertical graph)')
parser.add_argument("logFiles",nargs='*')
args=parser.parse_args()
NOF=len(args.logFiles)
allData=[]
ticks={}
try:
    for i in args.tick.split(','):
        ticks[i[0:2]]=i[2:]
except AttributeError:
    pass

if args.size:
    figureSize=(float(args.size.split(',')[0]),float(args.size.split(',')[1]))
else: 
    figureSize=None
for i in xrange(NOF):
    logFile=args.logFiles[i]
    if not os.path.isfile(logFile):
        print logFile+" do not exist"
        sys.exit()
    scanData=dihedralScan(logFile)
    scanData.unit=args.unit
    scanData.readLog()
    allData.append(scanData)
##find the y max of allData
ymaxmax=0 
for scanData in allData:
    ymax=max(scanData.energy)
    if ymax>ymaxmax:
        ymaxmax=ymax
if args.ymax:
    ymax=float(args.ymax)
else:
    ymax=ymaxmax*1.1

for i in xrange(NOF):  
    logFile=args.logFiles[i]
    logFileName=''.join(re.split('[/.]',logFile)[-2:-1]) 
        
    pylab.figure(1,figsize=figureSize)
    pylab.subplot(1,NOF,i+1)
    pylab.xlabel("Torsion angle (degree)")
    pylab.ylabel("Energy ("+allData[i].unit+")")
    pylab.title(logFileName)
    pylab.ylim(ymax=ymax)
    if 'hx' in ticks:
        pylab.xticks(pylab.arange(min(allData[i].coords),max(allData[i].coords),float(ticks['hx'])))
    if 'hy' in ticks:
        pylab.yticks(pylab.arange(min(allData[i].energy),max(allData[i].energy),float(ticks['hy'])))
    pylab.plot(allData[i].coords,allData[i].energy)
    
    pylab.figure(2,figsize=figureSize)
    pylab.subplot(NOF,1,i+1)
    pylab.xlabel("Torsion angle (degree)")
    pylab.ylabel("Energy ("+allData[i].unit+")")
    pylab.ylim(ymax=ymax)
    pylab.title(logFileName)
    if 'vx' in ticks:
        pylab.xticks(pylab.arange(min(allData[i].coords),max(allData[i].coords),float(ticks['vx'])))
    if 'vy' in ticks:
        pylab.yticks(pylab.arange(min(allData[i].energy),max(allData[i].energy),float(ticks['vy'])))
    pylab.plot(allData[i].coords,allData[i].energy)
    pylab.figure(3,figsize=figureSize)
    ma=i%3+3
    mb=i/3
    mc=0
    if ma==3 and mb==1:
        mb=3
    if mb>=3:
        mc=mb/3*45
        mb=mb%3
    pylab.plot(allData[i].coords,allData[i].energy,markersize=12,marker=(ma,mb,mc),label=logFileName)
    if 'ox' in ticks:
        pylab.xticks(pylab.arange(min(allData[i].coords),max(allData[i].coords),float(ticks['ox'])))
    if 'oy' in ticks:
        pylab.yticks(pylab.arange(min(allData[i].energy),max(allData[i].energy),float(ticks['oy'])))
    pylab.xlabel("Torsion angle (degree)")
    pylab.ylabel("Energy ("+scanData.unit+")")
    pylab.ylim(ymax=ymax)
    pylab.legend()
pylab.figure(2)
pylab.subplots_adjust(hspace=0.5)

if args.pic:
    for i in args.pic.split(','):
        if i in ['pdf','png','ps']:        
            pylab.figure(3).savefig(logFileName+'_torsionOne.'+i,dpi=args.dpi,format=i)
        elif i in ['eps','svg']:    
            pylab.figure(3).savefig(logFileName+'_torsionOne.'+i,dpi=args.dpi,format=i)
        if NOF > 1:
            pylab.figure(1).savefig(logFileName+'_torsionH.'+i,dpi=args.dpi,format=i)
            pylab.figure(2).savefig(logFileName+'_torsionV.'+i,dpi=args.dpi,format=i)
else:
    pylab.show()
##print test.energyKCal
